package demo.concurrent.executors.completionService;

import org.springframework.util.Assert;

import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Future;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;
import java.util.function.Consumer;
import java.util.function.Supplier;

/**
 * 等待缓存
 */
public class WaitCache {

    /** core */
    private static final Integer core = Runtime.getRuntime().availableProcessors();

    /** max */
    private static final Integer max = Runtime.getRuntime().availableProcessors() * 2;

    /** 锁映射 */
    private final Map<Object, ReentrantLock> lockMap = new ConcurrentHashMap<>();

    /** 用于锁住创建分段锁的地方 */
    private final Lock obj = new ReentrantLock();

    /** Executor service */
    private final ExecutorService executorService;

    /**
     * Wait cache
     *
     * @param executorService executor service
     * @since 1.0.0
     */
    public WaitCache(ExecutorService executorService) {
        this.executorService = Optional.ofNullable(executorService).orElseGet(() -> new ThreadPoolExecutor(core,
                                                                                                           max,
                                                                                                           60,
                                                                                                           TimeUnit.SECONDS,
                                                                                                           new ArrayBlockingQueue<>(500)));
    }

    /**
     * Wait cache
     *
     * @since 1.0.0
     */
    public WaitCache() {
        this(null);
    }

    /**
     * Write
     *
     * @param <T>       parameter
     * @param key       key
     * @param source    source
     * @param getSource get source
     * @param consumer  consumer
     * @since 1.0.0
     */
    public <T> void write(Object key, Supplier<T> source, Supplier<T> getSource, Consumer<T> consumer) {
        T t = source.get();
        retry:
        if (t == null) {
            Lock lock = createLock(key);
            lock.lock();
            try {
                Future<T> task = executorService.submit(source::get);
                t = task.get(200, TimeUnit.MILLISECONDS);
                if (t != null) {
                    break retry;
                } else {
                    if (getSource != null) {
                        t = executorService.submit(getSource::get).get(200, TimeUnit.MILLISECONDS);
                    } else {
                        return;
                    }
                }
                System.out.println(Thread.currentThread().getName() + "拿到缓存锁，设置数据：" + t);
            } catch (InterruptedException | ExecutionException | TimeoutException e) {
                System.out.println("获取数据超时，线程ID:" + Thread.currentThread().getName());
                e.printStackTrace();
                return;
            } finally {
                lock.unlock();
            }
        }
        consumer.accept(t);
    }

    /**
     * Create lock
     *
     * @param key key
     * @return the lock
     * @since 1.0.0
     */
    private Lock createLock(Object key) {
        ReentrantLock reentrantLock = lockMap.get(key);
        if (reentrantLock == null) {
            obj.lock();
            try {
                reentrantLock = Optional.ofNullable(lockMap.get(key)).orElseGet(() -> {
                    ReentrantLock lock = new ReentrantLock();
                    lockMap.put(key, lock);
                    return lock;
                });
            } finally {
                obj.unlock();
            }
        }
        return reentrantLock;
    }

    /**
     * 用于删除锁，注意如果删除了，可能会重复创建 ReentrantLock 导致效率低下
     *
     * @param key key
     * @since 1.0.0
     */
    public void release(Object key) {
        Assert.notNull(key, () -> "key不能为空");
        this.lockMap.remove(key);
    }

}
